package main

import "fmt"

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

func averageOfLevels(root *TreeNode) []float64 {
	if root == nil {
		return nil
	}
	queue := []*TreeNode{
		root,
	}
	depthMap := make(map[*TreeNode]int)
	depthMap[root] = 0
	var levels [][]int
	for len(queue) > 0 {
		head := queue[0]
		copy(queue, queue[1:])
		queue = queue[:len(queue)-1]
		depth := depthMap[head]
		if len(levels) <= depth {
			// first node
			levels = append(levels, []int{head.Val})
		} else {
			levels[depth] = append(levels[depth], head.Val)
		}
		if head.Left != nil {
			depthMap[head.Left] = depth + 1
			queue = append(queue, head.Left)
		}
		if head.Right != nil {
			depthMap[head.Right] = depth + 1
			queue = append(queue, head.Right)
		}
	}
	result := make([]float64, len(levels))
	for i, level := range levels {
		var sum float64
		for _, v := range level {
			sum += float64(v)
		}
		result[i] = sum / float64(len(level))
	}
	return result
}

func main() {
	root := &TreeNode{
		Val: 3,
		Left: &TreeNode{
			Val: 9,
		},
		Right: &TreeNode{
			Val: 20,
			Left: &TreeNode{
				Val: 15,
			},
			Right: &TreeNode{
				Val: 7,
			},
		},
	}
	ret := averageOfLevels(root)
	fmt.Printf("%#v\n", ret)
}
